from js_grid import *
import numpy as N
import math as M
import matplotlib.pyplot as plt
from Match_b_col import * # Imported to obtain the vectors tvec and so on

from Pot import *

from equations import init_vec, du_func, dv_func

n_var_u=6               # number of variables propagated in the u direction
n_var_v=6               # number of variables propagated in the v direction
nvar=n_var_u+n_var_v

udim= 200 ; vdim= 200

umin= 60. ; vmin = -59.9
umax= 95.1 ; vmax = 48.0

# I could put init_vec_col in a separate file, but it needs to know the parameters of the diamond
def init_vec_col(u,v):

  v = u+vmin-umin
  u = umin
  
  a=a_exact(u,v)
  r=r_exact(u,v)

  a_v=a_v_exact(u,v)
  r_v=r_v_exact(u,v)
  a_from_u=a
  r_from_u=r

  a_u=a_u_exact(u,v)
  r_u=r_u_exact(u,v)
  a_from_v=a
  r_from_v=r

  phi_from_u=phi_exact(u,v)
  phi_from_v=phi_exact(u,v)
  phi_u     =phi_u_exact(u,v)
  phi_v     =phi_v_exact(u,v)

  result=N.array([ 
                   a_u,
                   r_u,
                   a_from_u,
                   r_from_u,
                   phi_u,
                   phi_from_u,
                   a_v,
                   r_v,
                   a_from_v,
                   r_from_v,
                   phi_v,
                   phi_from_v,
                 ],dtype=N.double)
  return(result)

###################################################################

my_grid_high_res=Grid(udim,vdim,umin,umin+vmax-vmin,vmin,vmax,n_var_u,n_var_v,du_func,dv_func,init_vec_col,init_vec) # For collisions : init_vec_col is charged of the initial conditions at right
#my_grid_high_res=Grid(udim,vdim,umin,umax,vmin,vmax,n_var_u,n_var_v,du_func,dv_func,init_vec,init_vec)

###################################################################

grid_one=my_grid_high_res.w_grid

# cmap='prism'
def display_plot(i):
  a=grid_one[i,:,:]
  #vmin=a.min()
  #vmax=a.max()
  plt.close('all')
  #plt.imshow(a,vmin=vmin,vmax=vmax, interpolation='nearest',origin='lower')
  #vecu = N.linspace(0,udim-1,udim)*(umax-umin)/(udim-1)+umin # no collision
  vecu = N.linspace(0,udim-1,udim)*N.fabs(vmax-vmin)/(udim-1)+umin # for collisions
  plt.xlabel('u')
  vecv = N.linspace(0,udim-1,udim)*(vmax-vmin)/(udim-1)+vmin
  plt.ylabel('v')
  plt.contour(vecu,vecv,zip(*a),30)
  plt.colorbar()
  plt.show()

def display_energies():
  phi   = grid_one[11,:,:]
  phi_u = grid_one[4,:,:]
  phi_v = grid_one[10,:,:]
  a     = grid_one[2,:,:]
  rho = N.zeros((udim,udim),dtype=N.double)
  p   = N.zeros((udim,udim),dtype=N.double)
  for i in N.arange(udim):
    for j in N.arange(udim):
      rho[i,j]   = phi_u[i,j]*phi_v[i,j]/(2.*M.exp(2.*a[i,j]))+pot(phi[i,j])
      p[i,j]     = phi_u[i,j]*phi_v[i,j]/(2.*M.exp(2.*a[i,j]))-pot(phi[i,j])
  vecu = N.linspace(0,udim-1,udim)*N.fabs(vmax-vmin)/(udim-1)+umin
  vecv = N.linspace(0,udim-1,udim)*(vmax-vmin)/(udim-1)+vmin
  plt.subplot(2,1,1)
  plt.title('Density of matter')
  plt.xlabel('u')
  plt.ylabel('v')
  plt.contour(vecu,vecv,zip(*rho),30)
  #plt.imshow(rho,vmin=rho.min(),vmax=rho.max(), interpolation='nearest',origin='lower',cmap=plt.cm.gray)
  plt.subplot(2,1,2)
  plt.title('Pression')
  plt.xlabel('u')
  plt.ylabel('v')
  plt.contour(vecu,vecv,zip(*p),30)
  #plt.imshow(p,vmin=p.min(),vmax=p.max(), interpolation='nearest',origin='lower',cmap=plt.cm.gray)
  plt.colorbar()
  plt.show()

def display_speeds():
  phi_u = grid_one[4,:,:]
  phi_v = grid_one[10,:,:]
  a     = grid_one[2,:,:]
  speed_u = N.zeros((udim,udim),dtype=N.double)
  speed_v = N.zeros((udim,udim),dtype=N.double)
  for i in N.arange(udim):
    for j in N.arange(udim):
      if (phi_u[i,j]<0.):
        if (phi_v[i,j]<0.):
          speed_u[i,j] = M.sqrt(phi_u[i,j]/phi_v[i,j])/M.exp(a[i,j])
          speed_v[i,j] = M.sqrt(phi_v[i,j]/phi_u[i,j])/M.exp(a[i,j])
  vecu = N.linspace(0,udim-1,udim)*N.fabs(vmax-vmin)/(udim-1)+umin
  vecv = N.linspace(0,udim-1,udim)*(vmax-vmin)/(udim-1)+vmin
  vmin=min(speed_u.min(),speed_v.min())
  vmax=max(speed_u.max(),speed_v.max())
  plt.subplot(2,1,1)
  plt.title('Speed on u')
  plt.xlabel('u')
  plt.ylabel('v')
  plt.contour(vecu,vecv,zip(*speed_u),30)
  #plt.imshow(speed_u,vmin=vmin,vmax=vmax, interpolation='nearest',origin='lower')
  plt.subplot(2,1,2)
  plt.title('Speed on v')
  plt.xlabel('u')
  plt.ylabel('v')
  plt.contour(vecu,vecv,zip(*speed_v),30)
  #plt.imshow(speed_v,vmin=vmin,vmax=vmax, interpolation='nearest',origin='lower')
  plt.colorbar()
  plt.show()

plt.close('all')
